import torch as pt
import math

def calculate_ray_sphere_intersection(rays_local, rot, rays_o, radius):
  '''
  Args:
    rays_dir_local: [..., N, 3]; direction of rays in local coordinate
    rot: relative rotations between target and source, [..., 3, 3] matrices
    rays_o: [..., 3], origin position of rays

    radius: [..., 1]; radius of sphere
  Returns:
    intersections: [..., N, 3]; intersection points
  '''
  rays_dir = rays_local @ rot.transpose(-1, -2)
  rays_dir = pt.nn.functional.normalize(rays_dir, dim=-1)

  # (rays_o + rays_dir * t - sphere_center) ^ 2 == radius ^ 2
  # At^2 + Bt + C = 0, solve t

  A2 = pt.linalg.norm(rays_dir, dim=-1, keepdim=True) ** 2 * 2.0
  B = 2.0 * pt.sum(rays_dir * rays_o[...,None,:], dim=-1, keepdim=True)
  C = pt.linalg.norm(rays_o, dim=-1, keepdim=True) ** 2 - radius ** 2

  det = B ** 2 - 2.0 * A2 * C[...,None]
  sqrt_det = pt.abs(det) ** 0.5
  t = (-B + sqrt_det) / A2

  intersections = rays_o[...,None,:] + t * rays_dir
  mask = pt.logical_and(det > 0.0, t > 0.0)[...,0]
  return intersections, mask

# right-hand opencv coordinate system
def convert_point3d_to_normalized_theta_phi(
  points3d, theta_half_range, phi_half_range):
  theta_phi = pt.empty_like(points3d[...,:2], dtype=pt.float32)
  theta_phi[...,0] = pt.atan2(points3d[...,2], points3d[...,0]) # tan(theta) = z / x
  c = pt.linalg.norm(points3d[...,[0,2]], dim=-1)
  theta_phi[...,1] = pt.atan2(points3d[...,1], c) # tan(phi) = y / c

  theta_phi[...,0][theta_phi[...,0] < -math.pi * 0.5] += math.pi * 2
  theta_phi[...,0] = (math.pi / 2 - theta_phi[...,0]) / theta_half_range
  theta_phi[...,1] = theta_phi[...,1] / phi_half_range
  return theta_phi

def bilinear_wrapper_4d(imgs, coords):
  '''Wrapper around bilinear sampling function, handles arbitrary input sizes.

  Args:
    imgs: [..., H_s, W_s, C] images to resample
    coords: [..., N, 2], source pixel locations from which to copy
  Returns:
    [..., N, C] images after bilinear sampling from input
  '''
  new_imgs_shape = [coords.shape[i]
    if coords.shape[i] != imgs.shape[i] else -1 for i in range(imgs.dim() - 3)]
  new_imgs_shape += [-1, -1, -1]
  imgs = imgs.expand(new_imgs_shape)

  # The bilinear sampling code only handles 4D input, so we'll need to reshape.
  init_dims = imgs.shape[:-3]
  end_dims_img = imgs.shape[-3:]
  end_dims_coords = [coords.shape[-2], 1, coords.shape[-1]]
  prod_init_dims = init_dims.numel()

  imgs = imgs.reshape([prod_init_dims] + list(end_dims_img))
  imgs = pt.einsum('...ijk->...kij', imgs)
  coords = coords.reshape([prod_init_dims] + list(end_dims_coords))
  imgs_sampled = pt.nn.functional.grid_sample(imgs, coords, align_corners=True)
  imgs_sampled = imgs_sampled.transpose(-2, -3)
  return imgs_sampled.view(init_dims + imgs_sampled.shape[-3:-1])

def bilinear_wrapper_5d(imgs, coords):
  '''Wrapper around bilinear sampling function.

  Args:
    imgs: [layers, 1, H_s, W_s, C] images to resample
    coords: [total_layers, batch, N, 3], source pixel locations from which to copy
  Returns:
    [total_layers, batch, N, C] images after bilinear sampling from input
  '''

  imgs = imgs.permute([1, 4, 0, 2, 3])
  coords = coords.permute([1, 0, 2, 3])

  new_imgs_shape = [coords.shape[0] // imgs.shape[0], -1, -1, -1, -1]
  imgs = imgs.expand(new_imgs_shape)

  coords = coords.unsqueeze(-2)
  imgs_sampled = pt.nn.functional.grid_sample(imgs, coords, align_corners=True)
  imgs_sampled = imgs_sampled.permute(2, 0, 3, 1, 4)[...,0]
  return imgs_sampled

def warp_imgs_msi(imgs, pixel_rays_local, rot, t, r,
  theta_half_range, phi_half_range, offset=None):
  '''Transforms input hemisphere imgs to corresponding planes.

  Args:
    imgs: [..., H_s, W_s, C]
    pixel_rays_local: [..., N, 3]; rays direction in target coordinate
    rot: relative rotation, [..., 3, 3] matrices
    t: [..., 3], translations from target to source camera
      point p from target to source is accomplished via rot * p + t
    r: [..., 1], hemisphere radius
    theta_half_range: theta angle range with z axis in radians
    phi_half_range: phi angle range with z axis in radians
  Returns:
    [..., N, C] images after bilinear sampling from input.
      Coordinates outside the image are sampled as 0
  '''
  intersections, mask = calculate_ray_sphere_intersection(pixel_rays_local, rot, t, r)
  pt.cuda.empty_cache()

  if offset is not None:
    # offset is used as offset after getting ray sphere intersections
    theta_phi_coord = convert_point3d_to_normalized_theta_phi(
      intersections, theta_half_range, phi_half_range) # theta, phi in (-1, 1)
    offset_s2t = bilinear_wrapper_4d(offset, theta_phi_coord) * (r[1] - r[0]) * 0.5
    offset_dir = pt.nn.functional.normalize(
      intersections - t.unsqueeze(-2), dim=-1)
    intersections += offset_dir * offset_s2t

  theta_phi_coord = convert_point3d_to_normalized_theta_phi(
    intersections, theta_half_range, phi_half_range) # theta, phi in (-1, 1)
  imgs_s2t = bilinear_wrapper_4d(imgs, theta_phi_coord)
  imgs_s2t[~mask] = 0
  pt.cuda.empty_cache()

  return imgs_s2t
